//=============================================================================
//
//   Copyright (C) 2020 All rights reserved.
//
//   Filename:		dimension.rs
//
//   Author: Wang Zhecheng - wangzhecheng@yeah.net
//
//   Date: 2020-09-27 00:22
//
//   Last Editors: Wang Zhecheng - wangzhecheng@yeah.net
//
//   Last modified:	2020-11-06 20:55
//
//   Description:
//
//=============================================================================
use crate::Real;
use nalgebra::{
    allocator::Allocator,
    dimension::{DimName, U1, U2, U3, U6},
    {DefaultAllocator, Matrix2, Matrix3, Matrix6, MatrixN, Vector3, Vector6, VectorN},
};

pub trait DimType: 'static + std::fmt::Debug + Clone + Copy + Default {
    type DimVoigt: DimName;
    type Dim: DimName;

    fn get_elastic_matrix(e: Real, v: Real) -> MatrixN<Real, Self::DimVoigt>
    where
        DefaultAllocator: Allocator<Real, Self::DimVoigt, Self::DimVoigt>;
    fn get_perm_matrix(perm_x: Real, perm_y: Real, perm_z: Real) -> MatrixN<Real, Self::Dim>
    where
        DefaultAllocator: Allocator<Real, Self::Dim, Self::Dim>;
    fn from_voigt(rhs: &VectorN<Real, Self::DimVoigt>) -> MatrixN<Real, Self::Dim>
    where
        DefaultAllocator:
            Allocator<Real, Self::Dim, Self::Dim> + Allocator<Real, Self::DimVoigt, U1>;
    fn to_voigt(rhs: &MatrixN<Real, Self::Dim>) -> VectorN<Real, Self::DimVoigt>
    where
        DefaultAllocator:
            Allocator<Real, Self::Dim, Self::Dim> + Allocator<Real, Self::DimVoigt, U1>;
}

pub trait TwoDimension: DimType {}

#[derive(Debug, Copy, Clone, PartialEq, Default, serde::Deserialize)]
pub struct PlaneStrain {}
impl TwoDimension for PlaneStrain {}
impl DimType for PlaneStrain {
    type DimVoigt = U3;
    type Dim = U2;
    #[inline]
    fn get_elastic_matrix(e: Real, v: Real) -> Matrix3<Real> {
        let k0: Real = v / (1. - v);
        let g: Real = (1. - 2. * v) / (1. - v) * 0.5;
        let fac: Real = e * (1. - v) / (1. + v) / (1. - 2. * v);

        let dee = Matrix3::<Real>::new(
            1., k0, 0., //
            k0, 1., 0., //
            0., 0., g,
        );

        dee * fac
    }
    #[inline]
    fn get_perm_matrix(perm_x: Real, perm_y: Real, _: Real) -> Matrix2<Real> {
        Matrix2::<Real>::new(
            perm_x, 0 as Real, //
            0 as Real, perm_y, //
        )
    }
    #[inline]
    fn from_voigt(rhs: &Vector3<Real>) -> Matrix2<Real> {
        Matrix2::<Real>::new(
            rhs[0], rhs[2], //
            rhs[2], rhs[1], //
        )
    }
    #[inline]
    fn to_voigt(rhs: &Matrix2<Real>) -> Vector3<Real> {
        Vector3::<Real>::new(rhs[(0, 0)], rhs[(1, 1)], rhs[(1, 0)])
    }
}

#[derive(Debug, Copy, Clone, PartialEq, Default, serde::Deserialize)]
pub struct PlaneStress {}
impl TwoDimension for PlaneStress {}
impl DimType for PlaneStress {
    type DimVoigt = U3;
    type Dim = U2;
    #[inline]
    fn get_elastic_matrix(e: Real, v: Real) -> Matrix3<Real> {
        let k0: Real = v;
        let g: Real = (1. - v) * 0.5;
        let fac: Real = e / (1. - v * v);

        let dee = Matrix3::<Real>::new(
            1., k0, 0., //
            k0, 1., 0., //
            0., 0., g,
        );

        dee * fac
    }
    #[inline]
    fn from_voigt(rhs: &Vector3<Real>) -> Matrix2<Real> {
        Matrix2::<Real>::new(
            rhs[0], rhs[2], //
            rhs[2], rhs[1], //
        )
    }
    #[inline]
    fn to_voigt(rhs: &Matrix2<Real>) -> Vector3<Real> {
        Vector3::<Real>::new(rhs[(0, 0)], rhs[(1, 1)], rhs[(1, 0)])
    }
    #[inline]
    fn get_perm_matrix(perm_x: Real, perm_y: Real, _: Real) -> Matrix2<Real> {
        Matrix2::<Real>::new(
            perm_x, 0 as Real, //
            0 as Real, perm_y, //
        )
    }
}

#[derive(Debug, Copy, Clone, PartialEq, Default, serde::Deserialize)]
pub struct ThreeDimension {}
impl DimType for ThreeDimension {
    type DimVoigt = U6;
    type Dim = U3;
    #[inline]
    fn get_elastic_matrix(e: Real, v: Real) -> Matrix6<Real> {
        let k0: Real = v;
        let g: Real = (1. - v) * 0.5;
        let fac: Real = e / (1. - v * v);
        let dee = Matrix6::<Real>::new(
            1., k0, 0., 0., 0., 0., //
            k0, 1., 0., 0., 0., 0., //
            0., 0., 1., 0., 0., 0., //
            0., 0., 0., g, 0., 0., //
            0., 0., 0., 0., g, 0., //
            0., 0., 0., 0., 0., g,
        );
        dee * fac
    }
    #[inline]
    fn from_voigt(rhs: &Vector6<Real>) -> Matrix3<Real> {
        Matrix3::<Real>::new(
            rhs[0], rhs[3], rhs[4], //
            rhs[3], rhs[1], rhs[5], //
            rhs[4], rhs[5], rhs[2],
        )
    }
    #[inline]
    fn to_voigt(rhs: &Matrix3<Real>) -> Vector6<Real> {
        Vector6::<Real>::new(
            rhs[(0, 0)],
            rhs[(1, 1)],
            rhs[(2, 2)],
            rhs[(0, 1)],
            rhs[(0, 2)],
            rhs[(1, 2)],
        )
    }
    #[inline]
    fn get_perm_matrix(perm_x: Real, perm_y: Real, perm_z: Real) -> Matrix3<Real> {
        Matrix3::<Real>::new(
            perm_x, 0 as Real, 0 as Real, //
            0 as Real, perm_y, 0 as Real, //
            0 as Real, 0 as Real, perm_z,
        )
    }
}

pub trait HasDimension {
    fn get_dim(&self) -> usize;
    fn get_topo_dim(&self) -> usize;
}

#[test]
fn plane_strain_test() {
    let raw_voigt = Vector3::<Real>::new(1f32, 2f32, 3f32);
    let tensor = PlaneStrain::from_voigt(&raw_voigt);
    assert_eq!(tensor[(0, 0)], 1f32);
    assert_eq!(tensor[(0, 1)], 3f32);
    assert_eq!(tensor[(1, 0)], 3f32);
    assert_eq!(tensor[(1, 1)], 2f32);
    let voigt = PlaneStrain::to_voigt(&tensor);
    assert_eq!(voigt[(0)], 1f32);
    assert_eq!(voigt[(1)], 2f32);
    assert_eq!(voigt[(2)], 3f32);
    let dee = PlaneStrain::get_elastic_matrix(10_000f32, 0.3f32);
    assert_eq!(dee[(0, 0)], 13461.54f32);
    assert_eq!(dee[(0, 1)], 5769.232f32);
    assert_eq!(dee[(0, 2)], 0f32);
    assert_eq!(dee[(1, 0)], 5769.232f32);
    assert_eq!(dee[(1, 1)], 13461.54f32);
    assert_eq!(dee[(1, 2)], 0f32);
    assert_eq!(dee[(2, 0)], 0f32);
    assert_eq!(dee[(2, 1)], 0f32);
    assert_eq!(dee[(2, 2)], 3846.154f32);
    // ┌                            ┐
    // │ 13461.54 5769.232        0 │
    // │ 5769.232 13461.54        0 │
    // │        0        0 3846.154 │
    // └                            ┘
}
#[test]
fn plane_stress_test() {
    let raw_voigt = Vector3::<Real>::new(1f32, 2f32, 3f32);
    let tensor = PlaneStress::from_voigt(&raw_voigt);
    assert_eq!(tensor[(0, 0)], 1f32);
    assert_eq!(tensor[(0, 1)], 3f32);
    assert_eq!(tensor[(1, 0)], 3f32);
    assert_eq!(tensor[(1, 1)], 2f32);
    let voigt = PlaneStress::to_voigt(&tensor);
    assert_eq!(voigt[(0)], 1f32);
    assert_eq!(voigt[(1)], 2f32);
    assert_eq!(voigt[(2)], 3f32);
    let dee = PlaneStress::get_elastic_matrix(10_000f32, 0.3f32);
    assert_eq!(dee[(0, 0)], 10989.012f32);
    assert_eq!(dee[(0, 1)], 3296.7036f32);
    assert_eq!(dee[(0, 2)], 0f32);
    assert_eq!(dee[(1, 0)], 3296.7036f32);
    assert_eq!(dee[(1, 1)], 10989.012f32);
    assert_eq!(dee[(1, 2)], 0f32);
    assert_eq!(dee[(2, 0)], 0f32);
    assert_eq!(dee[(2, 1)], 0f32);
    assert_eq!(dee[(2, 2)], 3846.154f32);
    // ┌                               ┐
    // │ 10989.012 3296.7036         0 │
    // │ 3296.7036 10989.012         0 │
    // │         0         0  3846.154 │
    // └                               ┘
}
#[test]
fn three_dimension_test() {
    let raw_voigt = Vector6::<Real>::new(1f32, 2f32, 3f32, 4f32, 5f32, 6f32);
    let tensor = ThreeDimension::from_voigt(&raw_voigt);
    assert_eq!(tensor[(0, 0)], 1f32);
    assert_eq!(tensor[(0, 1)], 4f32);
    assert_eq!(tensor[(0, 2)], 5f32);
    assert_eq!(tensor[(1, 0)], 4f32);
    assert_eq!(tensor[(1, 1)], 2f32);
    assert_eq!(tensor[(1, 2)], 6f32);
    assert_eq!(tensor[(2, 0)], 5f32);
    assert_eq!(tensor[(2, 1)], 6f32);
    assert_eq!(tensor[(2, 2)], 3f32);
    // ┌       ┐
    // │ 1 4 5 │
    // │ 4 2 6 │
    // │ 5 6 3 │
    // └       ┘
    let voigt = ThreeDimension::to_voigt(&tensor);
    assert_eq!(voigt[(0)], 1f32);
    assert_eq!(voigt[(1)], 2f32);
    assert_eq!(voigt[(2)], 3f32);
    assert_eq!(voigt[(3)], 4f32);
    assert_eq!(voigt[(4)], 5f32);
    assert_eq!(voigt[(5)], 6f32);
    let dee = ThreeDimension::get_elastic_matrix(10_000f32, 0.3f32);
    let func = |r: usize, c: usize| {
        let pair = (r, c);
        match pair {
            (1, 0) => dee[(r, c)] == 3296.7036 as Real,
            (0, 1) => dee[(r, c)] == 3296.7036 as Real,
            (x, y) if x != y => dee[(r, c)] == 0 as Real,
            (x, _) if x > 2 => dee[(r, c)] == 3846.154 as Real,
            _ => dee[(r, c)] == 10989.012 as Real,
        }
    };
    assert!(Matrix6::<bool>::from_fn(func).iter().any(|&x| x));
    // ┌                                                             ┐
    // │ 10989.012 3296.7036         0         0         0         0 │
    // │ 3296.7036 10989.012         0         0         0         0 │
    // │         0         0 10989.012         0         0         0 │
    // │         0         0         0  3846.154         0         0 │
    // │         0         0         0         0  3846.154         0 │
    // │         0         0         0         0         0  3846.154 │
    // └                                                             ┘
}
